package com.shardingsphere.example.interceptor;


import com.shardingsphere.example.annotation.ShardingTable;
import com.shardingsphere.example.sharding.ShardingContextHolder;
import org.apache.ibatis.executor.statement.StatementHandler;
import org.apache.ibatis.mapping.BoundSql;
import org.apache.ibatis.mapping.MappedStatement;
import org.apache.ibatis.plugin.*;
import org.apache.ibatis.reflection.DefaultReflectorFactory;
import org.apache.ibatis.reflection.MetaObject;
import org.apache.ibatis.reflection.SystemMetaObject;
import org.springframework.stereotype.Component;

import java.sql.Connection;
import java.util.Properties;

/**
 * MyBatis分表拦截器
 */
@Intercepts({
        @Signature(type = StatementHandler.class, method = "prepare", args = {Connection.class, Integer.class})
})
@Component
public class ShardingInterceptor implements Interceptor {

    @Override
    public Object intercept(Invocation invocation) throws Throwable {
        StatementHandler statementHandler = (StatementHandler) invocation.getTarget();
        MetaObject metaObject = MetaObject.forObject(statementHandler,
                SystemMetaObject.DEFAULT_OBJECT_FACTORY,
                SystemMetaObject.DEFAULT_OBJECT_WRAPPER_FACTORY,
                new DefaultReflectorFactory());

        // 获取MappedStatement
        MappedStatement mappedStatement = (MappedStatement) metaObject.getValue("delegate.mappedStatement");

        // 获取方法上的ShardingTable注解
        String id = mappedStatement.getId();
        String className = id.substring(0, id.lastIndexOf("."));
        String methodName = id.substring(id.lastIndexOf(".") + 1);

        try {
            Class<?> clazz = Class.forName(className);

            // 检查类上的注解
            ShardingTable classAnnotation = clazz.getAnnotation(ShardingTable.class);

            // 检查方法上的注解
            ShardingTable methodAnnotation = null;
            try {
                methodAnnotation = clazz.getMethod(methodName).getAnnotation(ShardingTable.class);
            } catch (NoSuchMethodException e) {
                // 方法不存在，忽略
            }

            ShardingTable shardingTable = methodAnnotation != null ? methodAnnotation : classAnnotation;

            if (shardingTable != null) {
                // 获取SQL
                BoundSql boundSql = statementHandler.getBoundSql();
                String sql = boundSql.getSql();

                // 替换逻辑表名为实际表名
                String actualTableName = ShardingContextHolder.getTableName();
                if (actualTableName != null) {
                    String logicTable = shardingTable.logicTable();
                    String newSql = sql.replaceAll("\\b" + logicTable + "\\b", actualTableName);

                    // 设置新的SQL
                    metaObject.setValue("delegate.boundSql.sql", newSql);
                }
            }
        } catch (ClassNotFoundException e) {
            // 类不存在，忽略
        }

        return invocation.proceed();
    }

    @Override
    public Object plugin(Object target) {
        return Plugin.wrap(target, this);
    }

    @Override
    public void setProperties(Properties properties) {
        // 可以设置一些属性
    }
}

